#!/bin/bash

# 脚本名称：wipe_disk.sh

# 功能说明：
# 本脚本用于清空磁盘，删除所有分区。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-img、qemu-nbd、parted 等）。
# 2. 根据用户输入判断设备是物理磁盘还是虚拟磁盘镜像。
# 3. 如果是虚拟磁盘镜像，则使用 qemu-nbd 挂载到一个空闲的 nbd 设备上。
# 4. 显示设备信息。
# 5. 检查并卸载所有已挂载的分区。
# 6. 删除所有分区。
# 7. 断开 NBD 设备，完成操作。

# 注意事项：
# - 需要 root 权限运行部分命令，如 qemu-nbd、parted。
# - 在执行磁盘操作前，建议备份重要数据，以防万一。

# 安装依赖包
install_dependencies() {
    local required=("qemu-img" "qemu-nbd" "partprobe")
    local missing=()

    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -eq 0 ]; then
        return 0
    fi

    echo "检测到缺少依赖命令: ${missing[*]}，尝试自动安装..."
    sudo apt-get update -qq
    echo "正在安装 qemu-utils parted..."
    sudo apt-get install -y -qq qemu-utils parted >/dev/null
}

# 全局变量
disk=""
nbd_device=""
target_device=""
virtual_disk_attached=false

# 清理函数：如果虚拟磁盘已挂载，则断开 nbd 设备
cleanup() {
    if $virtual_disk_attached && [ -n "$nbd_device" ]; then
        echo "清理：断开 nbd 设备 $nbd_device"
        sudo qemu-nbd -d "$nbd_device"
    fi
}
trap cleanup EXIT SIGINT SIGTERM

# 查找空闲的 nbd 设备
select_free_nbd() {
    for dev in /dev/nbd*; do
        if [ -z "$(lsblk -n "$dev" | awk 'NR>1')" ]; then
            echo "$dev"
            return 0
        fi
    done
    return 1
}

# 挂载虚拟磁盘
mount_virtual_disk() {
    sudo modprobe nbd max_part=8
    nbd_device=$(select_free_nbd)
    if [ -z "$nbd_device" ]; then
        echo "未找到可用的 nbd 设备，退出。"
        exit 1
    fi
    echo "将虚拟磁盘镜像挂载到 nbd 设备：$nbd_device"
    sudo qemu-nbd -c "$nbd_device" "$disk"
    sudo udevadm settle
    virtual_disk_attached=true
}

# 获取用户输入
get_disk_input() {
    if [ -n "$1" ]; then
        disk="$1"
    else
        echo "没有提供设备或镜像文件，现列出所有磁盘："
        lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG -e7
        read -rp "请输入设备路径（如 /dev/sda）或虚拟磁盘镜像文件路径: " disk
    fi
}

# 处理设备类型
process_disk() {
    if [ -b "$disk" ]; then
        echo "选择的物理磁盘: $disk"
        target_device="$disk"
    elif [ -f "$disk" ]; then
        echo "选择的虚拟磁盘镜像文件: $disk"
        mount_virtual_disk
        target_device="$nbd_device"
    else
        echo "错误： '$disk' 既不是块设备也不是普通文件。"
        exit 1
    fi
}

# 显示设备信息
display_disk_info() {
    echo "显示设备信息："
    lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG "$target_device"
}

# 卸载所有已挂载的分区
umount_partitions() {
    echo "检查并卸载 $target_device 上的所有已挂载分区..."
    local partitions
    partitions=$(lsblk -nr -o NAME,MOUNTPOINT "$target_device" | awk '$2!="" {print $1}')
    for part in $partitions; do
        local part_dev="/dev/$part"
        echo "卸载分区 $part_dev..."
        sudo umount "$part_dev" || { echo "卸载 $part_dev 失败，退出！"; exit 1; }
    done
}

# 删除所有分区
remove_partitions() {
    echo "使用 parted 删除 $target_device 上的所有分区..."
    local partitions
    partitions=$(sudo parted -s "$target_device" print | awk '/^[[:space:]]*[0-9]+/ {print $1}') || true
    for part in $partitions; do
        echo "删除分区 $part..."
        sudo parted -s "$target_device" rm "$part"
    done
}

# 主函数
main() {
    install_dependencies
    get_disk_input "$1"
    process_disk
    display_disk_info
    umount_partitions
    remove_partitions
}

main "$@"

